from torch_geometric.datasets import Planetoid
from torch_geometric.loader import NeighborLoader
import torch
import numpy  as np
import networkx as nx
import matplotlib.pyplot as plt

data = Planetoid('./data', name='Cora')[0]

loader_2 = NeighborLoader(
    data,
    num_neighbors=[-1],
    batch_size=1,
    input_nodes=data.n_id,
)
# 准备边数据
sampled_data_2 = next(iter(loader_2))
# sampled_data_2 输出格式：
# Data(x=[4, 1433], edge_index=[2, 3], y=[4], train_mask=[4], val_mask=[4], test_mask=[4], n_id=[4], batch_size=1)
edge_2 = np.array(sampled_data_2.edge_index).T
edge_2 = edge_2.tolist()
edge_2 = list(tuple(line) for line in edge_2)

# 画图展示
G_2 = nx.Graph()
G_2.add_edges_from(edge_2)
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
option = {'font_family': 'serif', 'font_size': '15', 'font_weight': 'semibold'}
nx.draw_networkx(G_2, node_size=400, **option)
plt.show()